import os
import requests
import time
from fastapi import APIRouter

from ai_img.task import query_task_status

os.environ["DASHSCOPE_API_KEY"] = "sk-989ed32f49284be49d87158b4949b003"

text2img_app = APIRouter()


# def query_task_status(task_id):
#     # 设置 API 密钥
#     api_key = os.getenv("DASHSCOPE_API_KEY")
#     if not api_key:
#         raise ValueError("环境变量 DASHSCOPE_API_KEY 未设置")
#
#     # 请求 URL
#     url = f"https://dashscope.aliyuncs.com/api/v1/tasks/{task_id}"
#
#     # 请求头
#     headers = {
#         "Authorization": f"Bearer {api_key}"
#     }
#
#     # 发送 GET 请求
#     response = requests.get(url, headers=headers)
#
#     # 检查响应状态码
#     if response.status_code == 200:
#         print("请求成功，响应数据：")
#         print(response.json())
#         return response
#     else:
#         print(f"请求失败，状态码：{response.status_code}")
#         print(f"错误信息：{response.text}")


@text2img_app.post("/text2image")
def generate_image(param: dict):
    # 设置 API 密钥
    api_key = os.getenv("DASHSCOPE_API_KEY")
    if not api_key:
        raise ValueError("环境变量 DASHSCOPE_API_KEY 未设置")

    # 请求 URL
    url = "https://dashscope.aliyuncs.com/api/v1/services/aigc/text2image/image-synthesis"

    # 请求头
    headers = {
        "X-DashScope-Async": "enable",
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json"
    }

    # 请求体
    data = {
        "model": "wanx2.1-t2i-turbo",
        "input": {
            "prompt": param.get('prompt'),
            "negative_prompt": param.get('negative_prompt')
        },
        "parameters": {
            "size": param.get('size'),
            "n": param.get('num_images')
        }
    }

    # 发送 POST 请求
    response = requests.post(url, headers=headers, json=data)

    # 检查响应状态码
    if response.status_code == 200:
        print("请求成功，响应数据：")
        response_data = response.json()
        print(response_data)

        # 获取任务 ID
        task_id = response_data['output']['task_id']
        print(f"任务 ID: {task_id}")

        # 定时查询任务状态
        while True:
            status_response = query_task_status(task_id,api_key)
            status_data = status_response.json()

            # 检查任务状态
            if status_data['output']['task_status'] == 'SUCCEEDED':
                print("任务完成")
                result_url = status_data['output']['results'][0]['url']
                print(f"生成的图片 URL: {result_url}")
                return {'img_url': result_url}

            print("任务仍在处理中，等待 3 秒后重试...")
            time.sleep(3)
    else:
        print(f"请求失败，状态码：{response.status_code}")
        print(f"错误信息：{response.text}")


if __name__ == "__main__":
    prompt = "一间有着精致窗户的花店，漂亮的木质门，摆放着花朵 , 中国画风格"
    negative_prompt = "不要油画"
    generate_image(prompt, negative_prompt, '1280*720', 1)
